package com.study.living.mediacodec;

import android.content.Context;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.os.Build;

import androidx.annotation.RequiresApi;

import java.io.IOException;
import java.nio.ByteBuffer;

public class SurfaceEncoder extends Thread {

    private static final String TAG = "ProjectionService";

    //nal类型，取该字节的中间6位，之后右移1位，得到nal类型
    //例如：0x26
    private static final int NAL_I = 19;
    private static final int NAL_B = 0;
    //例如：0x02
    private static final int NAL_P = 1;

    //例如：0x40
    private static final int NAL_VPS = 32;
    private static final int NAL_SPS = 33;
    private static final int NAL_PPS = 34;

    private CallBack callBack;

    private MediaCodec mediaCodec;

    //nal分隔符占用的字节数，计算帧类型时需要偏移byte
    //分割符0x00 00 00 01 或者是0x00 00 01
    private int byteOffset;

    //帧类型
    private int nalType;

    //是否在每个I帧之前，添加vps、sps、pps信息。如果录屏的数据本地保存为mp4，不需要添加；如果网络实时传输，如直播，只需要添加
    private boolean addVpsSpsPpsBeforeIFrame = true;

    //保存vps数据，网络传输中放在I帧之前
    private byte[] vps_sps_pps_buf;
    private final CaptureScreenDataFile dataFile;

    public SurfaceEncoder(Context context) {
        dataFile = new CaptureScreenDataFile(context, "surface_encoder");
        callBack = dataFile::save;
    }

    @RequiresApi(api = Build.VERSION_CODES.LOLLIPOP)
    public void startEncoder(int width, int height) {
        try {
            mediaCodec = MediaCodec.createEncoderByType("video/hevc");

            MediaFormat mediaFormat = MediaFormat.createVideoFormat("video/hevc", width, height);
            //设置编码的颜色格式，这里是通过Surface
            mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface);
            //1080P最好在5Mbps/5120Kbps到8Mbps/8192Kbps之间,因为低于5Mbps不够清晰,而大于8Mbps视频文件会过大，比如我们设置8Mbps,则是1024*1024*8
            mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, 1024 * 1024 * 8);
            //帧率
            mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, 30);
            //关键帧间隔时间，单位为秒，此处的意思是这个视频每两秒一个关键帧
            mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 2);
            //最后一个参数需要注意，标明配置的是编码器
            mediaCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);

            mediaCodec.start();
        } catch (IOException e) {
            e.printStackTrace();
        }
        start();
    }

    @Override
    public void run() {
        super.run();
        startEncode();
        close();
        dataFile.closeFile();
    }

    private void startEncode() {
        MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
        int index;
        while (true) {
            try {
                index = mediaCodec.dequeueOutputBuffer(bufferInfo, 10_000);
                if (index >= 0) {
                    ByteBuffer outputBuffer;
                    if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.LOLLIPOP) {
                        outputBuffer = mediaCodec.getOutputBuffer(index);
                    } else {
                        outputBuffer = mediaCodec.getOutputBuffers()[index];
                    }
                    if (outputBuffer == null){
                        continue;
                    }
                    onEncodedDataAvailable(outputBuffer, bufferInfo);
                    mediaCodec.releaseOutputBuffer(index, false);
                }
            } catch (Exception e) {
                e.printStackTrace();
                break;
            }
        }
    }

    private void onEncodedDataAvailable(ByteBuffer outputBuffer, MediaCodec.BufferInfo bufferInfo) {
        if (addVpsSpsPpsBeforeIFrame){
            addVpsSpsPpsBeforeIFrame(outputBuffer, bufferInfo);
        } else {
            byte[] bytes = new byte[bufferInfo.size];
            outputBuffer.get(bytes);
            callBack.onEncodedDataAvailable(bytes);
        }
    }

    private void addVpsSpsPpsBeforeIFrame(ByteBuffer outputBuffer, MediaCodec.BufferInfo bufferInfo){
        if (outputBuffer.get(2) == 0x01) {
            //分割符0x00 00 01
            byteOffset = 3;
        } else {
            //分割符0x00 00 00 01
            byteOffset = 4;
        }
        //nal类型，取该字节的中间6位，之后右移1位，得到nal类型
        nalType = (outputBuffer.get(byteOffset) & 0x7E) >> 1;

        if (nalType == NAL_VPS) {
            vps_sps_pps_buf = new byte[bufferInfo.size];
            outputBuffer.get(vps_sps_pps_buf);
        } else if (nalType == NAL_I) {
            final byte[] bytes = new byte[vps_sps_pps_buf.length + bufferInfo.size];
            System.arraycopy(vps_sps_pps_buf, 0, bytes, 0, vps_sps_pps_buf.length);
            outputBuffer.get(bytes, vps_sps_pps_buf.length, bufferInfo.size);
            callBack.onEncodedDataAvailable(bytes);
        } else {
            byte[] bytes = new byte[bufferInfo.size];
            outputBuffer.get(bytes);
            callBack.onEncodedDataAvailable(bytes);
        }
    }

    public void close(){
        if (mediaCodec != null){
            try {
                mediaCodec.stop();
                mediaCodec.release();
            } catch (Exception e) {
                e.printStackTrace();
            }
            mediaCodec = null;
        }
    }

    /**
     * 编码后的数据
     */
    public interface CallBack {
        /**
         * 编码后的数据可用时。子线程
         * @param bytes
         */
        void onEncodedDataAvailable(byte[] bytes);
    }
}
